import flet as ft
import cv2
import threading
import time
import os
import numpy as np
from datetime import datetime
import mss
import pyautogui
from PIL import Image, ImageTk
import io
import base64

class VideoRecorder:
    def __init__(self):
        self.is_recording = False
        self.cap = None
        self.out = None
        self.current_file = None
        self._lock = threading.Lock()
        self.recording_type = "camera"  # "camera" 或 "screen"
        self.preview_frame = None
        self.camera_error = False  # 添加摄像头错误状态标志
        
    def start_recording(self, recording_type="camera"):
        with self._lock:
            try:
                self.is_recording = True
                self.recording_type = recording_type
                
                if recording_type == "camera":
                    if self.cap is None:
                        self.cap = cv2.VideoCapture(0)
                        # 添加摄像头初始化检查
                        if not self.cap.isOpened():
                            self.camera_error = True
                            raise Exception("无法打开摄像头，请检查摄像头是否正确连接")
                        
                        # 尝试读取一帧来验证摄像头是否正常工作
                        ret, _ = self.cap.read()
                        if not ret:
                            self.camera_error = True
                            raise Exception("无法从摄像头读取画面")
                            
                    frame_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                    frame_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                else:  # screen
                    screen_size = pyautogui.size()
                    frame_width, frame_height = screen_size.width, screen_size.height
                
                fps = 20.0
                
                if not os.path.exists("recordings"):
                    os.makedirs("recordings")
                    
                filename = f"recordings/{'screen' if recording_type == 'screen' else 'camera'}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4"
                self.current_file = filename
                
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                self.out = cv2.VideoWriter(filename, fourcc, fps, (frame_width, frame_height))

            except Exception as e:
                self.is_recording = False
                if self.cap is not None:
                    self.cap.release()
                    self.cap = None
                raise e

    def capture_screen(self):
        with mss.mss() as sct:  # 在函数内部创建 mss 实例
            monitor = sct.monitors[1]  # 主显示器
            screenshot = sct.grab(monitor)
            img = Image.frombytes('RGB', screenshot.size, screenshot.rgb)
            frame = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
            return frame

    def get_preview_frame(self):
        if not self.is_recording:
            return None
            
        try:
            if self.recording_type == "camera" and self.cap:
                if self.camera_error:
                    return None
                    
                ret, frame = self.cap.read()
                if not ret:
                    self.camera_error = True
                    print("无法获取摄像头画面")
                    return None
                    
                # 调整预览尺寸
                frame = cv2.resize(frame, (640, 480))
                return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            elif self.recording_type == "screen":
                frame = self.capture_screen()
                # 调整预览尺寸
                frame = cv2.resize(frame, (640, 480))
                return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        except Exception as e:
            print(f"获取预览画面时出错: {e}")
            return None
        return None
        
    def stop_recording(self):
        with self._lock:
            self.is_recording = False
            if self.out is not None:
                self.out.release()
                self.out = None
            if self.cap is not None:
                self.cap.release()
                self.cap = None
    
    def __del__(self):
        try:
            if self.out is not None:
                self.out.release()
            if self.cap is not None:
                self.cap.release()
        except Exception as e:
            print(f"清理资源时出错: {e}")

def main(page: ft.Page):
    page.title = "视频录制软件"
    page.theme_mode = ft.ThemeMode.LIGHT
    page.window.width = 800  # 增加窗口宽度
    page.window.height = 800  # 增加窗口高度
    page.padding = 20
    
    recorder = VideoRecorder()
    recording_thread = None
    timer_thread = None
    preview_thread = None
    
    # 状态显示
    status_text = ft.Text("就绪", size=20, text_align=ft.TextAlign.CENTER)
    timer_text = ft.Text("00:00", size=30, text_align=ft.TextAlign.CENTER)
    
    # 修改预览窗口
    preview_container = ft.Stack([
        ft.Container(
            width=640,
            height=480,
            border=ft.border.all(2, ft.Colors.GREY_400),
            border_radius=10,
            padding=5,
            visible=True
        ),
        ft.Image(
            width=640,
            height=480,
            fit=ft.ImageFit.CONTAIN,
            visible=False
        )
    ])

    def update_preview():
        image_control = preview_container.controls[1]
        last_update_time = time.time()
        frame_interval = 1.0 / 30.0  # 30 FPS
        
        while recorder.is_recording:
            try:
                current_time = time.time()
                if current_time - last_update_time < frame_interval:
                    time.sleep(0.001)  # 小睡避免CPU过度使用
                    continue
                    
                frame = recorder.get_preview_frame()
                if frame is not None:
                    # 将numpy数组转换为PIL图像
                    img = Image.fromarray(frame)
                    
                    # 创建临时内存缓冲区
                    img_byte_arr = io.BytesIO()
                    # 使用JPEG格式和较低的质量来提高性能
                    img.save(img_byte_arr, format='JPEG', quality=80)
                    img_byte_arr = img_byte_arr.getvalue()
                    
                    # 使用base64编码
                    img_base64 = base64.b64encode(img_byte_arr).decode()
                    
                    # 更新图像源
                    if not image_control.visible:
                        image_control.visible = True
                    image_control.src_base64 = img_base64
                    
                    # 使用batch更新来减少闪烁
                    page.update()
                    
                    last_update_time = current_time
                    
            except Exception as e:
                print(f"预览更新出错: {e}")
                time.sleep(0.1)  # 出错时稍微延长等待时间

    def record_video():
        while recorder.is_recording:
            try:
                with recorder._lock:
                    if recorder.is_recording:
                        if recorder.recording_type == "camera":
                            if recorder.cap and recorder.out:
                                ret, frame = recorder.cap.read()
                                if ret:
                                    recorder.out.write(frame)
                        else:  # screen
                            if recorder.out:
                                frame = recorder.capture_screen()
                                recorder.out.write(frame)
            except Exception as e:
                print(f"录制视频时出错: {e}")
                break

    def toggle_recording(recording_type):
        def handle_click(e):
            nonlocal recording_thread, timer_thread, preview_thread
            
            if not recorder.is_recording:
                try:
                    recorder.start_recording(recording_type)
                    status_text.value = "录制中..."
                    preview_container.visible = True
                    preview_container.controls[1].visible = False  # 重置图像控件状态
                    
                    # 隐藏录制按钮，显示停止按钮
                    buttons_container.content = ft.Row(
                        [stop_btn],
                        alignment=ft.MainAxisAlignment.CENTER,
                    )
                    
                    recording_thread = threading.Thread(target=record_video, daemon=True)
                    timer_thread = threading.Thread(target=update_timer, daemon=True)
                    preview_thread = threading.Thread(target=update_preview, daemon=True)
                    
                    recording_thread.start()
                    timer_thread.start()
                    preview_thread.start()
                except Exception as e:
                    error_message = str(e) if str(e) else "录制启动失败"
                    status_text.value = error_message
                    recorder.is_recording = False
                    preview_container.visible = False
                    # 如果出错，确保按钮状态正确
                    buttons_container.content = ft.Row(
                        [camera_btn, screen_btn],
                        alignment=ft.MainAxisAlignment.CENTER,
                    )
            
            page.update()
        return handle_click

    def stop_recording_handler(e):
        try:
            recorder.stop_recording()
            status_text.value = "视频已保存"
            timer_text.value = "00:00"
            preview_container.visible = False
            preview_container.controls[1].visible = False  # 重置图像控件状态
            
            # 显示录制按钮，隐藏停止按钮
            buttons_container.content = ft.Row(
                [camera_btn, screen_btn],
                alignment=ft.MainAxisAlignment.CENTER,
            )
            
            update_files_list()
        except Exception as e:
            print(f"停止录制时出错: {e}")
            status_text.value = "录制停止失败"
        
        page.update()

    def handle_window_event(e):
        if e.data == "close":
            page.open(confirm_dialog)

    def yes_click(e):
        try:
            # 停止录制
            if recorder.is_recording:
                recorder.stop_recording()
            
            # 强制设置录制状态为False
            recorder.is_recording = False
            
            # 关闭视频流
            if recorder.out is not None:
                recorder.out.release()
                recorder.out = None
            if recorder.cap is not None:
                recorder.cap.release()
                recorder.cap = None
            
            # 等待线程结束
            if recording_thread and recording_thread.is_alive():
                recording_thread.join(timeout=0.5)
            if timer_thread and timer_thread.is_alive():
                timer_thread.join(timeout=0.5)
            if preview_thread and preview_thread.is_alive():
                preview_thread.join(timeout=0.5)
            
            page.window.destroy()
        except Exception as e:
            print(f"关闭窗口时出错: {e}")
            page.window.destroy()

    def no_click(e):
        page.close(confirm_dialog)

    confirm_dialog = ft.AlertDialog(
        modal=True,
        title=ft.Text("确认退出"),
        content=ft.Text("确定要退出视频录制软件吗？"),
        actions=[
            ft.ElevatedButton("确定", on_click=yes_click),
            ft.OutlinedButton("取消", on_click=no_click),
        ],
        actions_alignment=ft.MainAxisAlignment.END,
    )

    # 创建按钮
    camera_btn = ft.ElevatedButton(
        "摄像头录制",
        icon=ft.Icons.VIDEOCAM,
        on_click=toggle_recording("camera"),
        style=ft.ButtonStyle(
            color=ft.Colors.WHITE,
            bgcolor=ft.Colors.BLUE_600,
            padding=20,
        ),
    )
    
    screen_btn = ft.ElevatedButton(
        "屏幕录制",
        icon=ft.Icons.DESKTOP_WINDOWS,
        on_click=toggle_recording("screen"),
        style=ft.ButtonStyle(
            color=ft.Colors.WHITE,
            bgcolor=ft.Colors.GREEN_600,
            padding=20,
        ),
    )
    
    stop_btn = ft.ElevatedButton(
        "停止录制",
        icon=ft.Icons.STOP,
        on_click=stop_recording_handler,
        style=ft.ButtonStyle(
            color=ft.Colors.WHITE,
            bgcolor=ft.Colors.RED_600,
            padding=20,
        ),
    )

    # 创建按钮容器
    buttons_container = ft.Container(
        content=ft.Row(
            [camera_btn, screen_btn],
            alignment=ft.MainAxisAlignment.CENTER,
        )
    )

    # 创建录制文件列表
    files_list = ft.ListView(
        expand=True,
        spacing=10,
        padding=20,
        height=200
    )
    
    def update_files_list():
        try:
            files_list.controls.clear()
            if os.path.exists("recordings"):
                for file in sorted(os.listdir("recordings"), reverse=True):
                    if file.endswith(".mp4"):
                        files_list.controls.append(
                            ft.Card(
                                content=ft.Container(
                                    content=ft.Text(file),
                                    padding=10
                                )
                            )
                        )
            page.update()
        except Exception as e:
            print(f"更新文件列表时出错: {e}")
    
    # 更新文件列表按钮
    refresh_btn = ft.IconButton(
        icon=ft.Icons.REFRESH,
        on_click=lambda _: update_files_list()
    )
    
    # 修改布局
    page.add(
        ft.Column(
            [
                ft.Container(
                    content=status_text,
                    alignment=ft.alignment.center,
                    padding=20
                ),
                ft.Container(
                    content=timer_text,
                    alignment=ft.alignment.center,
                    padding=20
                ),
                preview_container,  # 使用新的预览容器
                buttons_container,
                ft.Divider(height=40, thickness=2),
                ft.Row(
                    [
                        ft.Text("录制文件列表", size=16, weight=ft.FontWeight.BOLD),
                        refresh_btn
                    ],
                    alignment=ft.MainAxisAlignment.SPACE_BETWEEN
                ),
                files_list
            ],
            spacing=10,
            alignment=ft.MainAxisAlignment.START
        )
    )
    
    # 初始化文件列表
    update_files_list()
    
    # 设置窗口关闭事件处理
    page.window.prevent_close = True
    page.window.on_event = handle_window_event

    def update_timer():
        start_time = time.time()
        while recorder.is_recording:
            try:
                elapsed_time = int(time.time() - start_time)
                minutes = elapsed_time // 60
                seconds = elapsed_time % 60
                timer_text.value = f"{minutes:02d}:{seconds:02d}"
                page.update()
                time.sleep(1)
            except Exception as e:
                print(f"更新计时器时出错: {e}")
                break

if __name__ == "__main__":
    ft.app(target=main)